"""
掩模创建和管理工具

主要功能:
1. 交互式创建掩模区域
2. 支持多个掩模的创建和管理
3. 实时预览掩模效果
4. 支持撤销和重新绘制
"""

import cv2
import numpy as np
import my_tkinter
import tkinter
from tkinter import messagebox, filedialog
import time

# 全局变量定义
global img, k
k = 0  # 点击计数器
global mask_index
mask_index = 0  # 掩模索引
global point1, point2
lsPointsChoose = []  # 存储选择的点坐标(用于多边形ROI)
tpPointsChoose = []  # 存储选择的点坐标(用于绘制)
global first_point
first_point = (-10, -10)  # 第一个点的坐标
first_point_list = [(-10, -10), (-10, -10), (-10, -10), (-10, -10)]
tpPointsChoose_list = []  # 存储所有掩模的点坐标

mask_list = []  # 存储所有创建的掩模

def on_mouse(event, x, y, flags, param):
    """
    鼠标事件回调函数
    
    功能:
    1. 处理鼠标左键点击事件,记录点坐标
    2. 处理鼠标移动事件,实时预览掩模
    3. 处理鼠标右键点击事件,清除当前绘制
    4. 处理鼠标左键双击事件,完成当前掩模绘制
    """
    global img, point1, point2, count
    global lsPointsChoose, tpPointsChoose
    global init_img
    init_img = img.copy()  # 复制原图用于绘制

    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        global k
        k += 1
        if (k == 1):
            global first_point
            first_point = (x, y)
        point1 = (x, y)
        # 绘制点击位置
        cv2.circle(init_img, point1, 2, (0, 255, 0), 2)
        # 保存点坐标
        lsPointsChoose.append([x, y])
        tpPointsChoose.append((x, y))
        # 连接相邻点
        for i in range(len(tpPointsChoose) - 1):
            cv2.line(init_img, tpPointsChoose[i], tpPointsChoose[i + 1], (0, 0, 255), 2)
        cv2.imshow('src', init_img)
        
        # 检查是否完成一个闭合区域
        global distance
        if (distance < 4 and k > 1):
            global mask_index
            mask_index += 1
            tpPointsChoose_list.append(tpPointsChoose)
            tpPointsChoose = []
            vision_tpPointsChoose = []
            # 询问是否继续创建掩模
            continue_ = my_tkinter.Show_window('Do you want to continue to mask ? ')
            if (continue_ == 'y'):
                get_mask()
                tpPointsChoose = []
                vision_tpPointsChoose = []
                k = 0
                first_point = (-10, -10)
            else:
                get_mask()
                ROI_byMouse()
                root = tkinter.Tk()
                root.withdraw()
                messagebox.showinfo('提示', '点击图片窗口外缘，之后按任意键继续！')
                root.destroy()

    if (event == cv2.EVENT_MOUSEMOVE):  # 鼠标移动
        current_point = (x, y)
        # 计算当前点到第一个点的距离
        distance = (current_point[0] - first_point[0]) ** 2 + (current_point[1] - first_point[1]) ** 2
        if (distance < 4 and k > 1):
            # 如果接近起点,显示可闭合提示
            cv2.circle(init_img, current_point, 20, (0, 255, 0), 2)
            
        # 实时预览掩模轮廓
        vision_tpPointsChoose = []
        vision_tpPointsChoose = tpPointsChoose[::]
        vision_tpPointsChoose.append(current_point)
        for i in range(len(vision_tpPointsChoose) - 1):
            cv2.line(init_img, vision_tpPointsChoose[i], vision_tpPointsChoose[i + 1], (0, 0, 255), 2)
        cv2.imshow('src', init_img)
        for i in range(len(tpPointsChoose) - 1):
            cv2.line(init_img, tpPointsChoose[i], tpPointsChoose[i + 1], (0, 0, 255), 2)
        cv2.imshow('src', init_img)

    if event == cv2.EVENT_RBUTTONDOWN:  # 右键点击
        # 清除当前绘制
        tpPointsChoose = []
        k = 0
        first_point = (-10, -10)
        for i in range(len(tpPointsChoose) - 1):
            cv2.line(init_img, tpPointsChoose[i], tpPointsChoose[i + 1], (0, 0, 255), 5)
        cv2.imshow('src', init_img)

    if (event == cv2.EVENT_LBUTTONDBLCLK):  # 左键双击
        # 完成当前掩模绘制
        ROI_byMouse()

def get_mask():
    """
    根据选择的点创建掩模
    
    功能:
    1. 创建空白掩模图像
    2. 根据选择的点绘制多边形
    3. 填充多边形区域
    4. 将掩模添加到掩模列表
    """
    mask = np.zeros(img.shape, np.uint8)
    tpPointsChoose = tpPointsChoose_list[mask_index-1:mask_index]
    pts = np.array([tpPointsChoose], np.int32)
    pts = pts.reshape((-1, 1, 2))
    mask = cv2.polylines(mask, [pts], True, (255, 255, 255))
    mask = cv2.polylines(mask, [pts], True, 255)
    mask = cv2.fillPoly(mask, [pts], (255, 255, 255))
    mask_list.append(mask)

def ROI_byMouse():
    """
    应用掩模到原图
    
    功能:
    1. 遍历所有掩模
    2. 将掩模应用到原图
    3. 更新显示结果
    """
    global img
    for i in range(len(mask_list)):
        mask = mask_list[i]
        ROI = cv2.bitwise_and(~mask, img)
        img = ROI

def main(image_path=None):
    """
    主函数
    
    参数:
        image_path: 用于创建掩模的图像路径，如果为None则弹窗选择
    """
    global img, init_img, ROI, distance
    
    if image_path is None:
        # 如果没有提供图像路径，弹窗选择
        root = tkinter.Tk()
        root.withdraw()
        image_path = filedialog.askopenfilename(
            title='选择用于创建掩模的图像',
            filetypes=[('BMP files', '*.bmp'), ('All files', '*.*')]
        )
        if not image_path:
            raise ValueError("未选择图像文件")
    
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像文件: {image_path}")
        
    cv2.namedWindow('src')
    cv2.setMouseCallback('src', on_mouse)
    cv2.imshow('src', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return mask_list

if __name__ == '__main__':
    mask_list = main()